
import os

from PIL import Image
from lxml import etree
from xml.etree import ElementTree
from xml.etree.ElementTree import Element, SubElement


'''
将鸟窝(nw issue)和蜂巢(fc issue)修改为异物(yw issue)
'''


ENCODE_METHOD = 'utf-8'

base_dir = 'd:/image_augmentation/'
annotation_dir = base_dir + 'annotations/'


for p_dir,dirs,files in os.walk(annotation_dir):
    for name in files:
        xml_path = annotation_dir + name

        parser = etree.XMLParser(encoding=ENCODE_METHOD)
        tree = ElementTree.parse(xml_path, parser=parser)
        root = tree.getroot()

        #bndbox
        flag_save = False
        for object in root.iter('object'):
            catElement = object.find('name')
            cat = catElement.text

            if(cat == 'nw issue' or cat == 'fc issue'):
                catElement.text = 'yw issue'
                flag_save = True

        if(flag_save):
            tree.write(xml_path,encoding=ENCODE_METHOD,xml_declaration=True)

